package com.zf.experiment.limit;

import com.zf.experiment.annotation.InterceptorAccessLimit;
import com.zf.experiment.enums.AccessLimitMode;
import com.zf.experiment.exception.BusinessException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.lang.reflect.Method;
import java.util.Collections;
import java.util.List;
import java.util.Optional;

/**
 * 通过拦截器 使用 InterceptorAccessLimit 注解 限流
 *
 * @author Zhou Feng
 */
@Slf4j
@Component
public class InterceptorAccessLimitImpl implements HandlerInterceptor {

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (!(handler instanceof HandlerMethod)) {
            return true;
        }
        HandlerMethod handlerMethod = (HandlerMethod) handler;
        InterceptorAccessLimit accessLimit = handlerMethod.getMethodAnnotation(InterceptorAccessLimit.class);
        if (accessLimit == null || accessLimit.seconds() <= 0 || accessLimit.maxCount() < 0) {
            return true;
        }
        int seconds = accessLimit.seconds();
        int maxCount = accessLimit.maxCount();
        String key = getLimitKey(accessLimit, handlerMethod.getMethod());

        //读取 lua 脚本
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
        redisScript.setLocation(new ClassPathResource("lua/redis_limit.lua"));
        //必须设置
        redisScript.setResultType(Long.class);
        // redis中的key
        List<String> keys = Collections.singletonList(key);
        //执行 lua脚本，获取当前的访问数量
        Long tmpNumber = Optional.ofNullable(redisTemplate.execute(redisScript, keys, maxCount, seconds)).orElse(0L);
        log.info("{}", tmpNumber);
        //判断是否需要限流
        if (maxCount < tmpNumber.intValue()) {
            throw new BusinessException("限流了.....");
        }
        log.info("没有限流");
        return true;
    }

    public String getLimitKey(InterceptorAccessLimit accessLimit, Method method) {
        StringBuilder sb = new StringBuilder(accessLimit.prefix());
        // 按IP限流
        if (accessLimit.mode() == AccessLimitMode.IP) {
            ServletRequestAttributes requestAttributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
            String remoteAddr = requestAttributes.getRequest().getRemoteAddr();
            sb.append(remoteAddr).append(":");
        }
        // 按 访问数量限流
        Class<?> declaringClass = method.getDeclaringClass();
        sb.append(declaringClass.getName()).append("-").append(method.getName());
        return sb.toString();
    }
}
